import matplotlib.pyplot as plt
from parameters import *
import numpy as np


def line(w, b, x):
    print("true_k: ", -A / B)
    print("true_b: ", -C / B)
    print("predict_k: ", (1 / w[0][1]) * (- w[0][0]))
    print("predict_b: ", (1 / w[0][1]) * (- b[0]))
    return (1 / w[0][1]) * (- w[0][0] * x - b[0])


def draw(net, data_x, data_y):
    md = net.state_dict()
    print(md)
    W_tensor = md['hidden.weight']
    b_tensor = md['hidden.bias']
    W_np = W_tensor.numpy()
    b_np = b_tensor.numpy()
    X = np.linspace(-1, 1, 1000)
    plt.figure()
    plt.xlim((-1, 1))
    plt.ylim((-1, 1))
    plt.xlabel('x')
    plt.ylabel('y')
    plt.scatter(data_x[0], data_x[1], c=data_y)
    plt.plot(data_x[0], - A / B * data_x[0] - C / B, color='red', label='true')
    plt.plot(X, line(W_np, b_np, X), color="blue", label='predict')
    plt.legend()
    plt.show()
